// RVVAggregatorUDAF.java
package com.abcd1234;

import org.apache.spark.sql.expressions.Aggregator;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.Encoder;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder;

public class RVVAggregatorUDAF extends Aggregator<Row, RVVAggregatorUDAF.AggBuffer, Double[]> {
    private static final int BATCH_SIZE = 1024; // RVV向量处理批次大小

    public static class AggBuffer {
        long statePtr; // 本地聚合状态指针
        double[] quantities = new double[BATCH_SIZE];
        double[] prices = new double[BATCH_SIZE];
        double[] discounts = new double[BATCH_SIZE];
        double[] taxes = new double[BATCH_SIZE];
        int count = 0; // 当前批次中的元素数量
    }

    private final TPCHRVVAggregator nativeAgg = new TPCHRVVAggregator();

    @Override
    public AggBuffer zero() {
        AggBuffer buffer = new AggBuffer();
        buffer.statePtr = nativeAgg.createAggState();
        return buffer;
    }

    @Override
    public AggBuffer reduce(AggBuffer buffer, Row row) {
        // 将数据添加到缓冲区
        buffer.quantities[buffer.count] = row.getDouble(0); // l_quantity
        buffer.prices[buffer.count] = row.getDouble(1);     // l_extendedprice
        buffer.discounts[buffer.count] = row.getDouble(2);  // l_discount
        buffer.taxes[buffer.count] = row.getDouble(3);      // l_tax
        buffer.count++;
        
        // 当缓冲区满时，调用RVV优化函数
        if (buffer.count == BATCH_SIZE) {
            flushBuffer(buffer);
        }
        
        return buffer;
    }

    private void flushBuffer(AggBuffer buffer) {
        if (buffer.count > 0) {
            nativeAgg.updateBatch(
                buffer.statePtr, 
                buffer.quantities, buffer.prices, 
                buffer.discounts, buffer.taxes, 
                buffer.count
            );
            buffer.count = 0;
        }
    }

    @Override
    public AggBuffer merge(AggBuffer buffer1, AggBuffer buffer2) {
        // 先刷新缓冲区
        flushBuffer(buffer1);
        flushBuffer(buffer2);
        
        // 合并状态
        nativeAgg.merge(buffer1.statePtr, buffer2.statePtr);
        
        // 释放buffer2的状态
        nativeAgg.freeState(buffer2.statePtr);
        
        return buffer1;
    }

    @Override
    public Double[] finish(AggBuffer buffer) {
        // 刷新剩余数据
        flushBuffer(buffer);
        
        // 获取结果
        Double[] result = nativeAgg.getResult(buffer.statePtr);
        
        // 释放状态
        nativeAgg.freeState(buffer.statePtr);
        
        return result;
    }

    @Override
    public Encoder<AggBuffer> bufferEncoder() {
        return ExpressionEncoder.javaBean(AggBuffer.class);
    }

    @Override
    public Encoder<Double[]> outputEncoder() {
        return Encoders.bean(Double[].class);
    }
}
